package rate_limit

import (
	"context"

	"code.byted.org/bytertc/whiteboard_server/biz/config"
	"code.byted.org/gopkg/logs"
	"code.byted.org/kv/goredis"
	redis "code.byted.org/kv/redis-v6"
	"github.com/pkg/errors"
	"golang.org/x/time/rate"
)

const (
	MethodCreateTask = "create_task"
)

type RateLimiter struct {
	client       *goredis.Client
	localLimiter map[string]*rate.Limiter
}

func NewRateLimiter(client *goredis.Client) *RateLimiter {
	return &RateLimiter{
		client:       client,
		localLimiter: make(map[string]*rate.Limiter),
	}
}

var counterScript = redis.NewScript(`    
	redis.replicate_commands()
    local key1 = KEYS[1]

    local val = redis.call('incr', key1)
    local ttl = redis.call('ttl', key1)

    local expire = ARGV[1]
    local times = ARGV[2]

    if val == 1 then
        redis.call('expire', key1, tonumber(expire))
    else
        if ttl == -1 then
            redis.call('expire', key1, tonumber(expire))
        end
    end

	if val > tonumber(times) then
        return 0
    end

    return 1
`)

func (r *RateLimiter) Limit(ctx context.Context, method, appID string) bool {
	key := "whiteboard_rate_limiter" + method + appID
	expire := int64(1)
	accessLimit := int64(0)
	if cconfig, ok := config.Settings().RateLimiterConfig[method]; ok {
		accessLimit = cconfig[appID]
		if accessLimit == 0 {
			accessLimit = cconfig["*"]
		}
	}
	if accessLimit == 0 {
		return false
	}
	limit, err := r.redisLimit(ctx, key, expire, accessLimit)
	if err != nil {
		logs.CtxError(ctx, "redis limit error:%s", err)
		return r.localLimit(ctx, key, int(expire), int(accessLimit))
	}
	return limit
}

func (r *RateLimiter) redisLimit(ctx context.Context, key string, expire, accessLimit int64) (bool, error) {
	res, err := counterScript.Run(r.client, []string{key}, expire, accessLimit).Result()
	if err != nil {
		return false, err
	}
	if i, ok := res.(int64); ok {
		switch i {
		case 0:
			return true, nil
		case 1:
			return false, nil
		}
	}
	logs.CtxError(ctx, "redis error,res:%v", res)
	return false, errors.New("redis return error")
}

func (r *RateLimiter) localLimit(ctx context.Context, key string, expire, accessLimit int) bool {
	ll, ok := r.localLimiter[key]
	if !ok {
		ll = rate.NewLimiter(rate.Limit(accessLimit), accessLimit)
		r.localLimiter[key] = ll
	}
	if ll.Burst() != accessLimit {
		ll.SetLimit(rate.Limit(accessLimit))
		ll.SetBurst(accessLimit)
	}
	return !ll.Allow()
}
